'''
随机裁剪
usage
python user_det_random_clip_sg.py --path  D:\data\231215安全带\train_flball\_add_imgs  --sizerange  0.1,0.8

'''


# import imp
import argparse

import cv2
import time
import json
import math
import numpy as np
from numpy import *
import random
from random import randint
# from common.utils import *
from user_util import mkdir,listdir,cv_imread,rect_off,imrect,cv_imwrite
import sys
import os

def read_xml_annotation(in_file):
    import xml.etree.ElementTree as ET
    if not os.path.exists(in_file):
        return None

    try:
        tree = ET.parse(open(in_file, mode='rb'))
    except Exception as e:
        return None
    root = tree.getroot()
    bndboxlist = []

    for object in root.findall('object'):  # 找到root节点下的所有country节点
        bndbox = object.find('bndbox')  # 子节点下节点rank的值

        xmin = int(float(bndbox.find('xmin').text))
        xmax = int(float(bndbox.find('xmax').text))
        ymin = int(float(bndbox.find('ymin').text))
        ymax = int(float(bndbox.find('ymax').text))
        name = object.findtext('name')
        # print(xmin,ymin,xmax,ymax)
        bndboxlist.append(([xmin, ymin, xmax, ymax], name))
        # print(bndboxlist)
    #if root.find('object') is None:
    #    return None
    #bndbox = root.find('object').find('bndbox')
    return bndboxlist

def xml_dumps(bndboxs, fn, img):
    fn1,ext = os.path.split(fn)
    if img is None:
        h, w = 100, 100
    else:
        h, w = img.shape[:2]
    objs = ''
    for rect,name in bndboxs:
        xmin,ymin,xmax,ymax = rect
        objs += f'''
  <object>
    <type>bndbox</type>
    <name>{name}</name>
    <pose>Unspecified</pose>
    <truncated>0</truncated>
    <difficult>0</difficult>
    <bndbox>
      <xmin>{int(xmin)}</xmin>
      <ymin>{int(ymin)}</ymin>
      <xmax>{int(xmax)}</xmax>
      <ymax>{int(ymax)}</ymax>
    </bndbox>
  </object>
'''
    return f'''
    <annotation verified="no">
  <folder>.</folder>
  <filename>{fn1}</filename>
  <path>{fn}</path>
  <source>
    <database>Unknown</database>
  </source>
  <size>
    <width>{w}</width>
    <height>{h}</height>
    <depth>1</depth>
  </size>
  <segmented>0</segmented>
  {objs}
</annotation>
'''


def detect_jb1(img, labels):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    ss = cv2.Sobel(gray, cv2.CV_16S, 1, 0)
    ss = np.abs(ss)/2
    ss = ss.astype('uint8')
    thd, ss = cv2.threshold(ss, 10, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    k = 512
    blur1 = cv2.boxFilter(ss, 5, (k, k), normalize=True)
    min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(blur1)
    x, y = max_loc
    crop_w = k
    crop_h = k
    rec = [int(x - crop_w / 2), int(y - crop_h / 2), int(x + crop_w / 2), int(y + crop_h / 2)]
    img = cv2.cvtColor(ss, cv2.COLOR_GRAY2BGR)
    drawrect(img, rec)
    return img

def cal_iou(box1, box2):
    """
    :param box1: = [xmin1, ymin1, xmax1, ymax1]
    :param box2: = [xmin2, ymin2, xmax2, ymax2]
    :return: 
    """
    xmin1, ymin1, xmax1, ymax1 = box1
    xmin2, ymin2, xmax2, ymax2 = box2
    # 计算每个矩形的面积
    s1 = (xmax1 - xmin1) * (ymax1 - ymin1)  # C的面积
    s2 = (xmax2 - xmin2) * (ymax2 - ymin2)  # G的面积
 
    # 计算相交矩形
    xmin = max(xmin1, xmin2)
    ymin = max(ymin1, ymin2)
    xmax = min(xmax1, xmax2)
    ymax = min(ymax1, ymax2)
 
    w = max(0, xmax - xmin)
    h = max(0, ymax - ymin)
    area = w * h  # C∩G的面积
    iou = area / (s1 + s2 - area)
    irect = [xmin, ymin, xmax, ymax]
    return iou, area, s1, s2, irect

import random

def rand_rect(img, h, w):
    x = random.randint(0, img.shape[1]-w)
    y = random.randint(0, img.shape[0]-h)
    return [x, y, x+w, y+h]

def cal_max_iou(img, box0, labels, thd):
    outboxs = []
    x, y = box0[:2] # 裁剪框
    #box1 = extend_rect1(img, box0, -0.001)
    for i in range(len(labels)): # 遍历标签框
        box,name = labels[i]
        iou, area, s1, s2, irect = cal_iou(box0, box) # 计算裁剪框和标注框的交集
        ss = area/(s2+1) # 交集占原标签框的面积

        if ss >= thd: # 小于重合阈值不考虑该标签
            box11 = rect_off(irect, [-x, -y]) # 相对裁剪框的交集标签坐标， 交集不会在裁剪框之外
            outboxs.append([box11, name])

    return outboxs

def detect_jb(img, labels, pos_cnt = 1, sizerange=[0.1, 0.2]):
    imglabs = []
    h = min(img.shape[0], img.shape[1])
    #k = 1024
    i = 0
    a0 = int(sizerange[0]*h)
    a1 = int(sizerange[1]*h)
    for j in range(1000): # 裁剪n次
        if i>=pos_cnt:
            break
        #x, y = (rect[0]+rect[2])/2, (rect[1]+rect[3])/2
        #rec = [int(x - crop_w / 2), int(y - crop_h / 2), int(x + crop_w / 2), int(y + crop_h / 2)]
        k = random.randint(a0, a1)
        rec = rand_rect(img, k, k)
        outboxs = cal_max_iou(img, rec, labels, 0.7) # 单次 计算交集   thd：交集占原标签框的面积
        if len(outboxs) > 0 and len(outboxs) == len(labels): # 避免多个标签时可能删除低于thd的标签 user  20231120
            im = imrect(img, rec)
            imglabs.append([im, outboxs])
            i += 1

    return imglabs

def get_img_xml(img_fn):
    fn1, ext = os.path.splitext(img_fn)
    xml_fn = f'{fn1}.xml'
    if not os.path.exists(xml_fn):
        return None, None
        
    labels = read_xml_annotation(xml_fn)
    if labels is None or len(labels)==0:
        return None, None
    
    src = cv_imread(img_fn, 1)
    return src, labels

def read_yolo_txt(txt_fn, names, h=1, w=1):
    txtFile = open(txt_fn)
    # 读取所有内容
    txtList = txtFile.readlines()
    out = []
    for line in txtList:
        aa = [float(x) for x in line.split()]
        label = names[int(aa[0])]
        x_center = w * float(aa[1])       # aa[1]左上点的x坐标  
        y_center = h * float(aa[2])       # aa[2]左上点的y坐标
        width = (w*float(aa[3]))/2      # aa[3]图片width
        height = (h*float(aa[4]))/2      # aa[4]图片height
        lefttopx = int(x_center-width/2.0)
        lefttopy = int(y_center-height/2.0)
        #roi = img[lefttopy+1:lefttopy+height+3,lefttopx+1:lefttopx+width+1]
        rect = [x_center-width, y_center-height, x_center+width, y_center+height]
        rect = [round(x) for x in rect]
        out.append([rect, label])
        
    return out

def get_img_txt(img_fn, names):
    fn1, ext = os.path.splitext(img_fn)
    txt_fn = f'{fn1}.txt'
    if not os.path.exists(txt_fn):
        return None, None

    if not os.path.exists(img_fn):
        return None, None

    src = cv_imread(img_fn, 1)
    h, w = src.shape[:2]
    labels = read_yolo_txt(txt_fn, names, h, w)
    if labels is None or len(labels)==0:
        return None, None
    
    return src, labels




def ttest_cutimg(pa, sizerange=[0.1, 0.2], pos_cnt = 4):
    # fnlist = fnlist[]
    os.chdir(pa)
    # mkdir(f'{pa}/out')
    mkdir(f'{pa}_mini')
    # os.system('del out\\*.jpg')
    # os.system('del out\\*.xml')
    li = listdir('../', ['.jpg', '.png'])
    for i in range(len(li)): # 遍历图片
        img_fn = li[i]
        img_fn = f'{img_fn}'
        fn1, ext = os.path.splitext(img_fn)
        if ext not in ['.jpg', '.jpeg', '.png', '.bmp']:
            continue
        
        src, labels = get_img_xml(img_fn) # 读取图片和标签, 需要同时有图片和标签
        #src, labels = get_img_txt(img_fn)
        if src is None:  # 同时有图片和标签
            continue
        print(i, fn1)
        # pos_cnt = 4
        imglabs = detect_jb(src, labels, pos_cnt, sizerange) # 裁剪为pos_cnt个图片和标签

        if isinstance(imglabs, list):
            for i in range(len(imglabs)): # 保存图片和标签
                im,bndboxs = imglabs[i]
                cv_imwrite(f'{pa}_mini/{fn1}_{i}.jpg', im)
                open(f'{pa}_mini/{fn1}_{i}.xml', 'w').write(xml_dumps(bndboxs, f'{fn1}_{i}.jpg', im))
        else:
            cv_imwrite(f'{pa}_mini/{fn1}.jpg', imglabs)

'''
随机裁剪，
保存目录：{图片父文件夹}_mini
'''
def clip_muti_dirs(pa, sizerange, pos_cnt =4):
    # pa, sizerange = r'D:\data\231215安全带\train_flball\_add_imgs', [0.2, 0.8]

    # 获取子文件夹
    exts = ['.xml',]
    ls_dir = [] # 图片父文件夹, （含xml文件的文件夹~，且~不含_mini, 且同级目录不含~_mini）
    for root, dirs, files in os.walk(pa):  # os.walk 递归遍历
        for fn in files:
            fn0, ext = os.path.splitext(fn)
            isNewDir = not os.path.exists(root+'_mini') # 目录未裁剪过
            isOriginDir = '_mini' not in root # 目录非裁剪目录
            isXmlDir = ext.lower() in exts # 目录含xml文件
            if isNewDir and isOriginDir and isXmlDir: #
                ls_dir.append(root)
                break
    print(ls_dir)
    for i in ls_dir:
        ttest_cutimg(i, sizerange, pos_cnt)  # 一个文件夹内jpg png （不含子文件夹）


if __name__ == "__main__":
    # #pa, names = 'D:/data/221115保压/train_baoya', 'ok1 ok2 ng'
    #
    # # pa, sizerange = r'D:\data\231207霍尼\train_all\_add_imgs\20231207_simu', [0.4, 0.8]
    #
    # pa, sizerange = r'D:\data\220303九段冷却管\train_sunshang\_add_imgs\20240122', [0.2, 0.8]
    # #test_yolotxt2xml(pa, names.split())
    # #print('end')
    # test_cutimg(pa, sizerange, 4) # 一个文件夹内jpg png （不含子文件夹）
    # #pa, fun = '%s/out' % pa, detect_blob
    # #test_cutimg(pa, fun)

    parser = argparse.ArgumentParser()
    parser.add_argument('--path', type=str, default=r'D:\data\231215安全带\train_flball\_add_imgs', help='源目录')
    parser.add_argument('--sizerange', type=str, default='0.1,0.8', help='sizerange ')
    par = parser.parse_args()
    size = list(map(float, par.sizerange.split(',')))
    print(size)
    clip_muti_dirs(par.path, size)

    # clip_muti_dirs(r'D:\data\231215安全带\train_flball\_add_imgs', [0.1,0.8])